# langchain & langgraph
from langchain_core.tools import tool
import json

import jieba
import joblib
from langchain_core.tools import tool
import requests

from src.config import company_knn, company_tfidf, company_json
from src.choice_api.utils import filter_choose

stopwords = set(["有限公司", "嘉兴市", "浙江"])


def chinese_tokenizer(text):
    words = jieba.cut(text)  # 精确模式分词
    filtered = [word.strip() for word in words if word.strip() not in stopwords]
    return filtered


model = joblib.load(company_knn)
tfidf = joblib.load(company_tfidf)




# @tool("get_id")
def get_id(company_name):
    """
    根据公司名称查找对应id
    :param company_name:
    :return:
    """
    if company_name:
        with open(company_json, 'r', encoding='utf-8') as f:
            json_data = json.load(f)
            entities = list(json_data.keys())[0:-1]
            entities.extend(json_data[list(json_data.keys())[-1]]['entity'])
        query_text = [company_name]
        tfidf.tokenizer = chinese_tokenizer
        x_tfidf = tfidf.transform(query_text)
        distances, indices = model.kneighbors(x_tfidf)
        pos_company_name = entities[indices[0][0]]
        return {"info": json_data[pos_company_name]}
    return None


# @tool("get_upland_water_info")
# def get_upland_water_info(start_date, end_date, company_id=None):
#     """
#     查询上游来水数据 其中参数可以为空，为空参数设置为None
#     :param company_id:company_id根据简称从列表中选取['德山化工（浙江）有限公司','浙江嘉化能源化工股份有限公司','嘉兴市固体废物处置有限责任公司','浙江传化合成材料股份有限公司','合盛硅业（嘉兴）有限公司','浙江多凌线带科技股份有限公司','三江化工有限公司','嘉兴华雯化工股份有限公司','浙江博舍新材料有限公司','浙江恒优化纤有限公司','浙江嘉利宁环境科技有限公司','浙江益弹新材料科技有限公司','浙江惠禾源环境科技有限公司','浙江嘉福新材料科技有限公司']
#     :param start_date:
#     :param end_date:
#     :return:
#     """
#     url = api_ip + "/effluent/pretty-list1"
#     data = {
#         "companyId": company_id,
#         "startDate": start_date,
#         "endDate": end_date
#     }
#     return {"urls": url, "datas": data}
#
#     # return {"url": url, "data": data, "headers": headers}


@tool("get_upland_water_info")
def get_upland_water_info(startTime, endTime, company_id=None, target=None):
    """
    查询上游来水数据
    :param company_id:公司名称['德山化工','浙江嘉化', '浙江嘉化能源化工', '嘉兴市固体废物处置', '浙江传化合成材料', '合盛硅业', '浙江多凌线带科技', '三江化工','嘉兴华雯化工', '浙江博舍新材料', '浙江恒优化纤', '浙江嘉利宁环境科技', '浙江益弹新材料科技', '浙江惠禾源环境科技', '浙江嘉福新材料科技', '嘉兴瑞华泰薄膜技术']
    :param startTime:
    :param endTime:
    :param target:需要查询的指标，指标列表如下["flow", "ph", "cod", "ss", "nh3n", "tp", "tn"], default:NONE
    :return:
    """
    ids = []
    if isinstance(company_id, str):
        company_id = [company_id]

    if company_id:
        for id in company_id:
            ids.append(get_id(id))

    if not target or len(target) == 0:
        target = []

    url = "/effluent/pretty-list1"
    # print(type(ids[0]["info"]))
    # "companyId": [json.loads(item["info"])["text"]["companyId"] for item in ids],
    data = {
        "companyId": [item["info"]["text"]["companyId"] for item in ids],
        "startDate": startTime,
        "endDate": endTime
    }


    print("target:", target)

    return {"urls": url, "datas": data, "targets": [target]}


@tool("get_downstream_water_info")
def get_downstream_water_info(startTime, endTime, target=[]):
    """
    查询下游排海数据
    :param startTime:
    :param endTime:
    :param target:需要查询的指标，指标列表如下["flow", "ph", "cod", "ss", "nh3n", "tp", "tn"], default:NONE
    :return:
    """
    url = "/effluent/pretty-list3"

    if not target or len(target) == 0:
        target = []

    data = {
        "startDate": startTime,
        "endDate": endTime
    }
    return {"urls": url, "datas": data, "targets": [target]}
    # return {"info": f"上游来水接口信息:请求类型:post,url:{url},参数:{data},请求体:body", "urls":url, "datas":data}


@tool("get_lab_internal_info")
def get_lab_internal_info(startTime, endTime, checkStatus=None, sampleId=None, target=[]):
    """
    查询实验室检测分析数据（正常，异常），其中参数可以为空，为空参数设置为None
    :param checkStatus:实验室检测指标状态["正常","异常"], default:None
    :param sampleId:实验室检测位置列表["排放池", "磁混凝","臭氧池", "曝气生物滤池", "反硝化滤池(B)","反硝化滤池(A)","二期调节池","高效沉淀池","二沉池(B)","二沉池(A)","生化池1(ASR-B)","生化池1(ASR-A)","缺氧池(B)","缺氧池(A)","AnGSB池(B)","AnGSB池(A)","高浓度调节池","低浓度调节池","园区总管"], default:None
    :param startTime:
    :param endTime:
    :param target:需要查询的指标，指标列表如下['硫酸根', 'SS', 'COD(+AgNO3)', '硬度', '硫化物', '酸碱度', '碱度', '氯离子', '色度', 'MLVSS', 'MLSS','污泥沉降比', 'TSS', 'ALK', 'VFA', '余氯', '生化需氧量', 'TDS', '硝氮', 'Cond', '粪大肠菌群', 'pH','TP', '氨氮', 'TN', 'TOC', 'COD'], default:NONE
    :return:
    """
    url = "/lab/internal/list"
    checkStatus_mapping = {
        "正常": 1,
        "异常": 2
    }

    sampleId_mapping = {
        "排放池": 22,
        "磁混凝": 19,
        "臭氧池": 18,
        "曝气生物滤池": 16,
        "反硝化滤池(B)": 15,
        "反硝化滤池(A)": 14,
        "二期调节池": 13,
        "高效沉淀池": 12,
        "二沉池(B)": 11,
        "二沉池(A)": 10,
        "生化池1(ASR-B)": 9,
        "生化池1(ASR-A)": 8,
        "缺氧池(B)": 7,
        "缺氧池(A)": 6,
        "AnGSB池(B)": 5,
        "AnGSB池(A)": 4,
        "高浓度调节池": 3,
        "低浓度调节池": 2,
        "园区总管": 1
    }

    checkStatus = filter_choose(checkStatus, checkStatus_mapping)
    sampleId = filter_choose(sampleId, sampleId_mapping)

    if not target:
        target = []

    data = {
        "checkStatus": checkStatus,
        "sampleId": sampleId,
        "startTime": startTime,
        "endTime": endTime
    }
    return {"urls": url, "datas": data, "targets": [target]}


@tool("get_device_check_absence")
def get_device_check_absence(startTime, endTime, shift=None, type1=None, appealStatus=None):
    """
    查询厂内巡检缺卡/运行缺卡数据
    :param shift:班次列表["A班组","B班组", "C班组", "D班组"], default:None
    :param type1:["2H类", "12H类"], default:None
    :param appealStatus:["未申诉", "申述中","申述成功", "申述失败", "申述驳回"], default:None
    :param startTime:
    :param endTime:
    :return:
    """
    url = "/device-check-absence/list"

    shift_mapping = {
        "A班组": "A",
        "B班组": "B",
        "C班组": "C",
        "D班组": "D"
    }

    type_mapping = {
        "2H类": 2,
        "12H类": 12
    }

    appealStatus_mapping = {
        "未申诉": 1,
        "申述中": 2,
        "申述成功": 3,
        "申述失败": 4,
        "申述驳回": 5
    }

    shift = filter_choose(shift, shift_mapping)
    type1 = filter_choose(type1, type_mapping)
    appealStatus = filter_choose(appealStatus, appealStatus_mapping)

    data = {
        "shift": shift,
        "type": type1,
        "appealStatus": appealStatus,
        "startTime": startTime,
        "endTime": endTime
    }
    return {"urls": url, "datas": data, "targets": [[]]}


@tool("get_device_check")
def get_device_check(startTime, endTime, repairState=None, abnormal=None):
    """
    查询厂内运行巡检(巡检打卡/运行打卡、巡检正常/巡检异常、巡检异常维修情况)数据
    :param repairState:厂内修理状态["无需维修","维修单生成", "维修已完成", "正在维修"], default:None
    :param abnormal:厂内是否异常["正常情况", "异常情况"], default:None
    :param startTime:
    :param endTime:
    :return:
    """

    repairState_mapping = {
        "无需维修": 0,
        "维修单生成": 1,
        "正在维修": 2,
        "维修已完成": 3
    }

    abnormal_mapping = {
        "异常情况": 1,
        "正常情况": 2
    }

    repairState = filter_choose(repairState, repairState_mapping)
    abnormal = filter_choose(abnormal, abnormal_mapping)

    url = "/check-order/device-check/list"
    data = {
        "repairState": repairState,
        "abnormal": abnormal,
        "startTime": startTime,
        "endTime": endTime
    }
    return {"urls": url, "datas": data, "targets": [[]]}


@tool("get_internal_online_data")
def get_internal_online_data(startTime, endTime, target=None):
    """
    查询港区污水厂厂内在线数据(TOC,TP,TN,COD,PH)
    :param startTime:
    :param endTime:
    :param target:需要查询的指标，指标列表如下["toc", "tp", "tn", "cod", "pH"], default:NONE
    :return:
    """
    url = "/check-order/device-check/list"
    if not target or len(target) == 0:
        target = []
    data = {
        "startTime": startTime,
        "endTime": endTime
    }
    return {"urls": url, "datas": data, "targets": [target]}


@tool("get_electric_consume_record")
def get_electric_consume_record(startTime, endTime, type1=None):
    """
    查询电量能耗数据
    :param startTime:
    :param endTime:
    :param type1:设备类型列表["1#机主","2#机主", "深度处理A", "深度处理B", "新风机A", "新风机B", "新风机C"], default:None
    :return:
    """
    url = "/electric-consume-record/list"

    type1_mapping = {
        "1#机主": 1,
        "2#机主": 2,
        "深度处理A": 3,
        "深度处理B": 4,
        "新风机A": 5,
        "新风机B": 6,
        "新风机C": 7
    }

    type1 = filter_choose(type1, type1_mapping)

    data = {
        "startTime": startTime,
        "endTime": endTime,
        "type": type1
    }
    return {"urls": url, "datas": data, "targets": [[]]}


@tool("get_composite_consume_record")
def get_composite_consume_record(startTime, endTime, type1=None):
    """
    查询用水量(工业用水、生活用水)能耗数据
    :param startTime:
    :param endTime:
    :param type1:用水类型列表["工业用水","生活用水"], default:None
    :return:
    """
    url = "/composite-consume-record/list"
    type1_mapping = {
        "工业用水": 1,
        "生活用水": 2
    }

    type1 = filter_choose(type1, type1_mapping)

    data = {
        "startTime": startTime,
        "endTime": endTime,
        "type": type1
    }
    return {"urls": url, "datas": data, "targets": [[]]}


@tool("get_steam_dosage_record")
def get_steam_dosage_record(startTime, endTime):
    """
    查询用蒸汽量能耗数据
    :param startTime:
    :param endTime:
    :return:
    """
    url = "/steam_dosage-record/list"
    data = {
        "startTime": startTime,
        "endTime": endTime,
    }
    return {"urls": url, "datas": data, "targets": [[]]}


@tool("get_drug_inventory")
def get_drug_inventory(startTime, endTime, name=None, type1=None, deliverer=None):
    """
    查询药剂管控(出库、入库、库存、供应商)数据
    :param startTime:
    :param endTime:
    :param name: 药品名称列表['液体柠檬酸', '次钠', '高效复合碳源', '乙酸钠', '次氯酸钠固体', '消泡剂', '液体尿素', '尿素', '磷酸氢二铵', 'PAM(S806)(阴)', 'PAM(4490)(阳)', 'PAC', '液碱', '液氧', '活性炭', '营养液', '硫酸亚铁'] default:None
    :param type1:进出库状态['出库','入库'], default:出入库
    :param deliverer:药剂供货商名称['信安贸易有限公司', '嘉兴市嘉赢环境工程有限公司', '浙江巨能环境工程有限公司', '嘉兴市浩晖贸易有限公司', '嘉兴市南湖区首创助剂厂', '上海协聪化工有限公司', '嘉兴市堃达化工有限公司', '林德（上海）工业气体有限公司', '浙江嘉通新材料科技有限公司', '江苏正力环保科技有限公司', '南京欣瑞新材料科技有限公司', '嘉兴市禾晟生物制品有限公司', '嘉兴市创伟新材料股份有限公司']default:None
    :return:
    """
    url = "/drug-inventory/list"

    name_mapping = {
        "液体柠檬酸": "液体柠檬酸",
        "次钠药剂": "次钠",
        "高效复合碳源": "高效复合碳源",
        "乙酸钠": "乙酸钠",
        "次氯酸钠固体": "次氯酸钠固体",
        "消泡剂": "消泡剂",
        "液体尿素": "液体尿素",
        "尿素": "尿素",
        "磷酸氢二铵": "磷酸氢二铵",
        "PAM(S806)(阴)": "PAM(S806)(阴)",
        "PAM(4490)(阳)": "PAM(4490)(阳)",
        "PAC": "PAC",
        "液碱": "液碱",
        "液氧": "液氧",
        "活性炭": "活性炭",
        "营养液": "营养液",
        "硫酸亚铁": "硫酸亚铁"
      }

    type1_mapping = {
        "出库": 2,
        "入库": 1
      }

    deliverer_mapping = {
        "信安贸易有限公司": "13",
        "嘉兴市嘉赢环境工程有限公司": "12",
        "浙江巨能环境工程有限公司": "11",
        "嘉兴市浩晖贸易有限公司": "10",
        "嘉兴市南湖区首创助剂厂": "9",
        "上海协聪化工有限公司": "8",
        "嘉兴市堃达化工有限公司": "7",
        "林德（上海）工业气体有限公司": "6",
        "浙江嘉通新材料科技有限公司": "5",
        "江苏正力环保科技有限公司": "4",
        "南京欣瑞新材料科技有限公司": "3",
        "嘉兴市禾晟生物制品有限公司": "2",
        "嘉兴市创伟新材料股份有限公司": "1"
      }

    name = filter_choose(name, name_mapping)
    type1 = filter_choose(type1, type1_mapping)
    deliverer = filter_choose(deliverer, deliverer_mapping)



    data = {
        "startTime": startTime,
        "endTime": endTime,
        "name": name,
        "type": type1,
        "deliverer": deliverer
    }
    return {"urls": url, "datas": data, "targets": [[]]}


@tool("get_sludge_transport_record")
def get_sludge_transport_record(startTime, endTime, timeType=None, destination=None):
    """
    按按进场/入库时间查询污泥外运数据以及污泥含水量数据
    :param startTime:
    :param endTime:
    :param destination:['嘉爱斯', '平湖石化'] default:None
    :param timeType: 是以入库时间查询还是以出库时间查询 ['出库','入库'],default:'出库'
    :return:
    """
    url = "/sludge-transport-record/list"

    destination_mapping = {
        "嘉爱斯": "嘉爱斯",
        "平湖石化": "平湖石化"
    }

    destination = filter_choose(destination, destination_mapping)

    if "入库" in str(timeType):
        data = {
            "timeInStart": startTime,
            "timeInEnd": endTime,
            "destination": destination
        }
        target = ["in"]
    else:
        data = {
            "timeOutStart": startTime,
            "timeOutEnd": endTime,
            "destination": destination
        }
        target = ["out"]

    return {"urls": url, "datas": data, "targets": [target]}


# @tool("get_sludge_transport_record_out")
# def get_sludge_transport_record_out(timeOutStart, timeOutEnd, destination=None):
#     """
#     按出场/出库时间查询污泥外运数据以及污泥含水量外运数据
#     :param timeOutStart:
#     :param timeOutEnd:
#     :param destination:['嘉爱斯', '平湖石化'] default:None
#     :return:
#     """
#     url = "/sludge-transport-record/list"
#     data = {
#         "timeOutStart": timeOutStart,
#         "timeOutEnd": timeOutEnd,
#         "destination": destination
#     }
#     return {"urls": url, "datas": data, "targets": [[]]}


@tool("get_staff_attendance")
def get_staff_attendance(startTime, endTime, timeResult=None):
    """
    查询人员考勤打卡数据
    :param startTime:
    :param endTime:
    :param timeResult: 打卡情况列表['已打卡', '缺卡/未打卡', '早退', '迟到', '严重迟到', '旷工迟到', '休息'] default:None
    :return:
    """
    url = "/attendance/list"
    timeResult_mapping = {
        "已打卡": "Normal",
        "缺卡/未打卡": "NotSigned",
        "早退": "Early",
        "迟到": "Late",
        "严重迟到": "SeriousLate",
        "旷工迟到": "Absenteeism",
        "休息": "Rest"
    }

    timeResult = filter_choose(timeResult, timeResult_mapping)

    data = {
        "startTime": startTime,
        "endTime": endTime,
        "timeResult": timeResult
    }
    return {"urls": url, "datas": data, "targets": [[]]}


@tool("get_device")
def get_device(state=None):
    """
    查询设备清单(运行中设备、维修中/故障设备、异常状态设备、停用设备)/设备详情数据
    :param state: 设备状态列表['运行中', '维修中/故障情况', '异常', '停用']
    :return:
    """

    state_mapping = {
        "运行中": 0,
        "维修中/故障情况": 1,
        "异常": 2,
        "停用": 3
    }

    state = filter_choose(state, state_mapping)

    url = "/device/list"
    data = {
        "state": state
    }
    return {"urls": url, "datas": data, "targets": [[]]}
